import torch
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops
from torch_geometric.data import Data


class DualUpdateLayer(MessagePassing):
    def __init__(self, node_dim, edge_dim):
        super().__init__(aggr='mean')
        # 节点更新网络  node_dim * 2 + edge_dim  ==》》 node_dim
        self.node_mlp = torch.nn.Sequential(
            torch.nn.Linear(node_dim * 2 + edge_dim, 64),
            torch.nn.ReLU(),
            torch.nn.Linear(64, node_dim)
        )
        # 边更新网络  node_dim * 2 + edge_dim ==》》 edge_dim
        self.edge_mlp = torch.nn.Sequential(
            torch.nn.Linear(node_dim * 2 + edge_dim, 64),
            torch.nn.ReLU(),
            torch.nn.Linear(64, edge_dim)
        )

    def forward(self, x, edge_index, edge_attr):
        # 添加自环保证节点自更新
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
        return self.propagate(edge_index, x=x, edge_attr=edge_attr)

    def message(self, x_i, x_j, edge_attr):
        # 拼接源节点、目标节点和边特征
        combined = torch.cat([x_i, x_j, edge_attr], dim=-1)
        return {'node_msg': combined, 'edge_msg': combined}

    def aggregate(self, inputs, index):
        # 标准均值聚合
        return super().aggregate(inputs['node_msg'], index)

    def update(self, aggr_out, inputs):
        # 节点特征更新
        new_x = self.node_mlp(aggr_out)
        # 边特征更新
        new_edge = self.edge_mlp(inputs['edge_msg'])
        return new_x, new_edge

    def edge_update(self, edge_attr, x_i, x_j):
        # 独立边更新接口
        return self.edge_mlp(torch.cat([x_i, x_j, edge_attr], dim=-1))


if __name__ == "__main__":
    # 测试数据：4个节点/4条边
    x = torch.randn(4, 16)  # 节点特征16维
    edge_index = torch.tensor([[0, 1, 2, 3], [1, 2, 3, 0]], dtype=torch.long)
    edge_attr = torch.randn(4, 8)  # 边特征8维

    # 初始化更新层
    updater = DualUpdateLayer(node_dim=16, edge_dim=8)

    # 执行更新
    new_x, new_edge = updater(x, edge_index, edge_attr)

    print("原始节点尺寸:", x.shape)
    print("更新后节点:", new_x.shape)
    print("原始边尺寸:", edge_attr.shape)
    print("更新后边:", new_edge.shape)
